import os
import shutil
import subprocess

# 定义路径
source_dir = "rust-pro"
output_dir = "../Output"
src_source_dir = os.path.join(source_dir, "src")
src_test_dir = "src-test"
src_test_main_rs = os.path.join(src_test_dir, "src", "main.rs")
tests_dir = os.path.join(output_dir, "tests")
common_dir = os.path.join(tests_dir, "common")

# 定义Rust重构工具的路径
rust_refactor_tool = os.path.join("rust_refactor", "target", "debug", "rust_refactor")

if not os.path.isfile(rust_refactor_tool):
    print(f"错误: 找不到Rust重构工具: {rust_refactor_tool}")
    sys.exit(1)

# 调用Rust重构工具重构source_dir目录
try:
    print(f"正在运行Rust重构工具: {rust_refactor_tool} {source_dir}")
    result = subprocess.run([rust_refactor_tool, source_dir], check=True, capture_output=True, text=True)
    print("Rust重构工具运行成功。输出如下:")
    print(result.stdout)
except subprocess.CalledProcessError as e:
    print("Rust重构工具运行失败。错误信息如下:")
    print(e.stderr)

# 如果 Output 目录不为空，则删除整个目录
if os.path.exists(output_dir) and os.listdir(output_dir):
    print("Output directory is not empty. Deleting all contents...")
    shutil.rmtree(output_dir)

# 创建 Output 目录
os.makedirs(output_dir, exist_ok=True)

# 在 Output 目录下创建 Cargo.toml 文件
cargo_toml_content = """\
[package]
name = "Output"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
libc = "0.2.155"
lazy_static = "1.4.0"
"""
with open(os.path.join(output_dir, "Cargo.toml"), "w") as f:
    f.write(cargo_toml_content)

# 用于记录哪些文件编译成功
successful_files = set()

# 遍历 rust-pro/src 目录下的所有 .rs 文件
for file_name in os.listdir(src_source_dir):
    if file_name.endswith(".rs") and file_name != "lib.rs":  # 排除 lib.rs 文件
        file_path = os.path.join(src_source_dir, file_name)
        with open(file_path, "r") as f:
            content = f.read()

        # 检查是否包含 fn main
        if "fn main" not in content:
            # 在文件末尾添加空的 fn main
            content += "\nfn main() {}\n"

        # 将修改后的内容写入 src-test/src/main.rs
        with open(src_test_main_rs, "w") as f:
            f.write(content)

        # 在 src-test 目录下运行 cargo build
        print(f"Testing {file_name} with cargo build...")
        result = subprocess.run(["cargo", "build"], cwd=src_test_dir)

        # 如果 cargo build 成功，继续处理
        if result.returncode == 0:
            print(f"cargo build succeeded for {file_name}. Copying to Output directory...")

            # 创建 Output/src 目录（如果不存在）
            os.makedirs(os.path.join(output_dir, "src"), exist_ok=True)

            # 拷贝原始文件到 Output/src 目录
            shutil.copy(file_path, os.path.join(output_dir, "src", file_name))

            # 更新 lib.rs 文件
            lib_rs_path = os.path.join(output_dir, "src", "lib.rs")
            module_name = file_name[:-3]  # 去掉 .rs 后缀
            with open(lib_rs_path, "a") as f:
                f.write(f"pub mod {module_name};\n")

            # 记录编译成功的文件
            successful_files.add(file_name)
        else:
            print(f"cargo build failed for {file_name}. Skipping...")

# 创建 tests 和 tests/common 目录
os.makedirs(common_dir, exist_ok=True)

# 复制 alloc_testing.rs 到 common 目录
alloc_testing_src = os.path.join(source_dir, "test", "alloc_testing.rs")
framework_src = os.path.join(source_dir, "test", "framework.rs")

if os.path.exists(alloc_testing_src):
    shutil.copy(alloc_testing_src, os.path.join(common_dir, "alloc_testing.rs"))

# 处理 framework.rs：添加对 alloc_testing.rs 的导入，及 null_test_function 定义
if os.path.exists(framework_src):
    with open(framework_src, "r") as f:
        framework_content = f.read()

    # 在第一行添加导入
    if "use super::alloc_testing" not in framework_content:
        framework_content = "use super::alloc_testing::{alloc_test_get_allocated, alloc_test_set_limit};\n" + framework_content

    # 如果没有 null_test_function，添加一个空的函数
    if "fn null_test_function" not in framework_content:
        framework_content += "\nfn null_test_function() {}\n"

    # 保存修改后的 framework.rs 文件
    with open(os.path.join(common_dir, "framework.rs"), "w") as f:
        f.write(framework_content)

# 创建 common/mod.rs 文件
mod_rs_content = """\
// common/mod.rs
pub mod alloc_testing;
pub mod framework;
"""
with open(os.path.join(common_dir, "mod.rs"), "w") as f:
    f.write(mod_rs_content)

# 根据编译结果拷贝特定的测试文件
hash_files = {"hash_int.rs", "hash_pointer.rs", "hash_string.rs"}
compare_files = {"compare_int.rs", "compare_pointer.rs", "compare_string.rs"}

# 检查是否有 hash 文件编译成功
if any(file in successful_files for file in hash_files):
    test_hash_functions_src = os.path.join(source_dir, "test", "test_hash_functions.rs")
    if os.path.exists(test_hash_functions_src):
        with open(test_hash_functions_src, "r") as f:
            hash_content = f.read()

        # 添加模块导入
        hash_content = """mod common;
use common::alloc_testing::*;
use common::framework::{null_test_function, run_tests, UnitTestFunction};
""" + hash_content

        # 检查 hash 文件并添加相关导入
        if "hash_int.rs" in successful_files:
            hash_content = "use Output::hash_int::int_hash;\n" + hash_content
        if "hash_pointer.rs" in successful_files:
            hash_content = "use Output::hash_pointer::pointer_hash;\n" + hash_content
        if "hash_string.rs" in successful_files:
            hash_content = "use Output::hash_string::{string_hash, string_nocase_hash};\n" + hash_content

        # 保存修改后的 test_hash_functions.rs 文件
        with open(os.path.join(tests_dir, "test_hash_functions.rs"), "w") as f:
            f.write(hash_content)

# 检查是否有 compare 文件编译成功
if any(file in successful_files for file in compare_files):
    test_compare_functions_src = os.path.join(source_dir, "test", "test_compare_functions.rs")
    if os.path.exists(test_compare_functions_src):
        with open(test_compare_functions_src, "r") as f:
            compare_content = f.read()

        # 添加模块导入
        compare_content = """mod common;
use common::alloc_testing::*;
use common::framework::{null_test_function, run_tests, UnitTestFunction};
""" + compare_content

        # 检查 compare 文件并添加相关导入
        if "compare_int.rs" in successful_files:
            compare_content = "use Output::compare_int::*;\n" + compare_content
        if "compare_pointer.rs" in successful_files:
            compare_content = "use Output::compare_pointer::*;\n" + compare_content
        if "compare_string.rs" in successful_files:
            compare_content = "use Output::compare_string::*;\n" + compare_content

        # 保存修改后的 test_compare_functions.rs 文件
        with open(os.path.join(tests_dir, "test_compare_functions.rs"), "w") as f:
            f.write(compare_content)

# 处理 test_alloc_testing.rs
test_alloc_testing_src = os.path.join(source_dir, "test", "test_alloc_testing.rs")
if os.path.exists(test_alloc_testing_src):
    with open(test_alloc_testing_src, "r") as f:
        alloc_testing_content = f.read()

    # 添加模块导入
    alloc_testing_content = """mod common;
use common::alloc_testing::*;
use common::framework::{null_test_function, run_tests, UnitTestFunction};
""" + alloc_testing_content

    # 保存修改后的 test_alloc_testing.rs 文件
    with open(os.path.join(tests_dir, "test_alloc_testing.rs"), "w") as f:
        f.write(alloc_testing_content)

# 要特殊处理的模块列表
special_modules = ["arraylist", "avl_tree", "binary_heap", "binomial_heap", "bloom_filter", "hash_table", "list", "rb_tree", "slist", "sorted_array", "queue", "tinyexpr", "trie", "utf8_decoder", "set"]
single_modules = ["queue", "tinyexpr", "trie", "utf8_decoder"]

# 处理特殊模块的测试文件
for module in special_modules:
    module_rs = f"{module}.rs"
    if module_rs in successful_files:
        test_module_src = os.path.join(source_dir, "test", f"test_{module}.rs")
        if os.path.exists(test_module_src):
            with open(test_module_src, "r") as f:
                module_content = f.read()

            # 创建导入语句
            import_statements = """mod common;
use common::alloc_testing::*;
use common::framework::\{null_test_function, run_tests, UnitTestFunction\};
"""         
            import_statements += f"use Output::{module}::*;\n"
            if module not in single_modules:
                # 针对 hash_table 模块，按单个文件编译结果添加相关导入
                if module == "hash_table":
                    # 定义需要检查的文件及其对应的导入语句
                    dependency_imports = {
                        "hash_int.rs": "use Output::hash_int::*;\n",
                        "compare_int.rs": "use Output::compare_int::*;\n",
                        "hash_string.rs": "use Output::hash_string::*;\n",
                        "compare_string.rs": "use Output::compare_string::*;\n"
                    }
                    # 遍历每个依赖文件，如果编译成功，则添加对应的导入语句
                    for dep_file, import_stmt in dependency_imports.items():
                        if dep_file in successful_files:
                            import_statements += import_stmt
                # 如果处理的是 bloom_filter，并且 hash_string 编译成功，则额外添加导入
                elif module == "set":
                # 定义需要检查的文件及其对应的导入语句
                    dependency_imports = {
                        "hash_int.rs": "use Output::hash_int::*;\n",
                        "compare_int.rs": "use Output::compare_int::*;\n",
                        "hash_pointer.rs": "use Output::hash_pointer::*;",
                        "compare_pointer.rs": "use Output::compare_pointer::*;",
                        "hash_string.rs": "use Output::hash_string::*;\n",
                        "compare_string.rs": "use Output::compare_string::*;\n"
                    }
                    # 遍历每个依赖文件，如果编译成功，则添加对应的导入语句
                    for dep_file, import_stmt in dependency_imports.items():
                        if dep_file in successful_files:
                            import_statements += import_stmt
                elif module == "bloom_filter" and "hash_string.rs" in successful_files:
                    import_statements += "use Output::hash_string::*;\n"
                # 检查 compare 文件并添加相关导入
                elif "compare_int.rs" in successful_files:
                    import_statements += "use Output::compare_int::*;\n" + compare_content
            # 将导入语句添加到文件头部
            module_content = import_statements + module_content

            # 保存修改后的测试文件
            with open(os.path.join(tests_dir, f"test_{module}.rs"), "w") as f:
                f.write(module_content)

            print(f"Updated test_{module}.rs with necessary imports.")

print("Output directory structure and files created successfully.")
